import os
import cv2
from MyTools import kitti_2_yolo
import random
import shutil
import sys
import xml.etree.ElementTree as ET
from tqdm import tqdm
from utils.general import download, Path, os
import fiftyone as fo
import fiftyone.zoo as foz
import fiftyone as fo
import fiftyone.zoo as foz
import fiftyone.zoo as foz
from fiftyone import ViewField as F




# def train_test(img_path, label_path, test_img, test_label):
#     img_files = os.listdir(img_path)
#     num_imgs = len(img_files)
#     num_to_move = int(num_imgs * 0.3)  # 计算需要移动的图片数量
    
#     os.makedirs(test_img, exist_ok=True)
#     os.makedirs(test_label, exist_ok=True)

#     # 随机选择30%的图片
#     selected_imgs = random.sample(img_files, num_to_move)
#     for img_file in selected_imgs:
#         # 移动图片文件
#         img_src = os.path.join(img_path, img_file)
#         img_dst = os.path.join(test_img, img_file)
#         shutil.move(img_src, img_dst)
        
#         # 移动对应的标签文件
#         label_file = os.path.splitext(img_file)[0] + '.txt' 
#         label_src = os.path.join(label_path, label_file)
#         label_dst = os.path.join(test_label, label_file)
#         if os.path.exists(label_src):  
#             shutil.move(label_src, label_dst)

def mkdir(name = 'kitti'):
    img_train_save_folder = f"../datasets/{name}/images/train"
    img_val_save_folder = f"../datasets/{name}/images/val"
    img_test_save_folder = f"../datasets/{name}/images/test"
    label_train_save_folder = f"../datasets/{name}/labels/train"
    label_val_save_folder = f"../datasets/{name}/labels/val"
    label_test_save_folder = f"../datasets/{name}/labels/test"
    
    try: 
        os.makedirs(img_train_save_folder)
    except:
        None
    try: 
        os.makedirs(img_val_save_folder)
    except:
        None
    try: 
        os.makedirs(img_test_save_folder)
    except:
        None
        
    try: 
        os.makedirs(label_train_save_folder)
    except:
        None
    try: 
        os.makedirs(label_val_save_folder)
    except:
        None
    try: 
        os.makedirs(label_test_save_folder)
    except:
        None
    return ([img_train_save_folder, img_val_save_folder, img_test_save_folder],
            [label_train_save_folder, label_val_save_folder, label_test_save_folder])

def cp_2_target(img_path, label_path, img_save_folder, label_save_folder, img_name):
    label_name = os.path.splitext(img_name)[0] + ".txt"
    train_img_path = os.path.join(img_path, img_name)
    train_label_path = os.path.join(label_path, label_name)
    
    if not os.path.exists(train_img_path) or not os.path.exists(train_label_path):
        sys.exit()

    shutil.copy(train_img_path, img_save_folder)
    shutil.copy(train_label_path, label_save_folder)


def split_dataset(img_path, label_path = "../converted_label/k_2_y", shuffle_img = True, 
                  split_percent = [0.56, 0.14, 0.3], 
                  name = 'kitti'):
    # (img_train_save_folder, img_val_save_folder, img_test_save_folder,
    #  label_train_save_folder, label_val_save_folder, label_test_save_folder) = mkdir(name)#先拿到保存路径
    img_folders, label_folders = mkdir(name)

    all_imgs = os.listdir(img_path) # 得到所有文件的名字
    if shuffle_img:
        #打乱所有图片
        random.shuffle(all_imgs)

    img_num = len(all_imgs)
    train_num = int(img_num * split_percent[0])
    val_num = int(img_num * split_percent[1])
    test_num = img_num - train_num - val_num
    
    train_imgs = all_imgs[ : train_num]
    val_imgs = all_imgs[train_num : train_num + val_num]
    test_imgs = all_imgs[train_num + val_num : ]
    imgs_sets = [train_imgs, val_imgs, test_imgs]

    #开始写入目标文件夹
    
    for imgs_set, img_folder, label_folder in zip(imgs_sets, img_folders, label_folders):
        for i, train_img in enumerate(imgs_set):
            if train_img.endswith(('.jpg', '.png', '.jpeg')):
                cp_2_target(img_path, 
                            label_path, 
                            img_folder, 
                            label_folder,  
                            train_img) #得到一个图片和label的路径
            if i % (train_num // 10) == 0:
                progress = (i / train_num) * 100
                print(f"已完成 {progress:.0f}% ({i}/{train_num})")
            

    # for i, val_img in enumerate(val_imgs):
    #     if val_img.endswith(('.jpg', '.png', '.jpeg')):
    #         cp_2_target(img_path, 
    #                     label_path, 
    #                     img_val_save_folder, 
    #                     label_val_save_folder,  
    #                     val_img) #得到一个图片和label的路径
    #     if i % (val_num // 10) == 0:
    #         progress = (i / val_num) * 100
    #         print(f"已完成 {progress:.0f}% ({i}/{val_num})")
            




VOC_name = {
    # 手动调配了。前8个是kitti的。
    # 有个比较麻烦的是大小写，kitti的是大写的，Voc是小写的。
    0: 'car',
    1: 'van',
    2: 'truck',
    3: 'tram',
    4: 'person',
    5: 'person_sitting',
    6: 'cyclist', # 骑车和自行车还是有区别的
    7: 'misc',

    # VOC独有的
    8: 'aeroplane',
    9: 'bicycle',
    10: 'bird',
    11: 'boat',
    12: 'bottle',
    13: 'bus',
    14: 'cat',
    15: 'chair',
    16: 'cow',
    17: 'diningtable',
    18: 'dog',
    19: 'horse',
    20: 'motorbike',
    21: 'pottedplant',
    22: 'sheep',
    23: 'sofa',
    24: 'train',
    25: 'tvmonitor',
}

def VOC_convert_label(path, lb_path, year, image_id):
  def convert_box(size, box):
      dw, dh = 1. / size[0], 1. / size[1]
      x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]
      return x * dw, y * dh, w * dw, h * dh

  in_file = open(path / f'VOC{year}/Annotations/{image_id}.xml')
  out_file = open(lb_path, 'w')
  tree = ET.parse(in_file)
  root = tree.getroot()
  size = root.find('size')
  w = int(size.find('width').text)
  h = int(size.find('height').text)

  names = list(VOC_name.values())  # names list
  for obj in root.iter('object'):
      cls = obj.find('name').text
      if cls in names and int(obj.find('difficult').text) != 1:
          xmlbox = obj.find('bndbox')
          bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')])
          cls_id = names.index(cls)  # class id
          out_file.write(" ".join([str(a) for a in (cls_id, *bb)]) + '\n')

def VOC_split(root, VOC2007, VOC2012):
    dir = Path(root)  # dataset root dir
    path = dir
    for year, image_set in ('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test'):
      imgs_path = dir / 'images' / f'{image_set}{year}'
      lbs_path = dir / 'labels' / f'{image_set}{year}'
      imgs_path.mkdir(exist_ok=True, parents=True)
      lbs_path.mkdir(exist_ok=True, parents=True)
    
      with open(path / f'VOC{year}/ImageSets/Main/{image_set}.txt') as f:
          image_ids = f.read().strip().split()
      for id in tqdm(image_ids, desc=f'{image_set}{year}'):
          f = path / f'VOC{year}/JPEGImages/{id}.jpg'  # old img path
          lb_path = (lbs_path / f.name).with_suffix('.txt')  # new label path
          #f.rename(imgs_path / f.name)  # move image
          shutil.copy(f, imgs_path / f.name)  # 复制文件，而不是移动
          VOC_convert_label(path, lb_path, year, id)  # convert labels to YOLO format



def download_openimages(
        max_samples=6000,
        seed=51,
        shuffle=True,
        down_load_class = ["Person", "Billboard", "Bicycle", "Bird", "Boat", "Rabbit", "Sheep", "Truck",
           "Van", "Monkey", "Pig", "Toy", "Traffic light", "Traffic sign", "Bus"]
    ):
    origin_dict = {
        0: "Car",
        1: "Van",
        2: "Truck",
        3: "Tram",
        4: "Person",
        5: "Person_sitting",
        6: "Cyclist",
        7: "Misc",
        8: "Aeroplane",
        9: "Bicycle",
        10: "Bird",
        11: "Boat",
        12: "Bottle",
        13: "Bus",
        14: "Cat",
        15: "Chair",
        16: "Cow",
        17: "Diningtable",
        18: "Dog",
        19: "Horse",
        20: "Motorbike",
        21: "Pottedplant",
        22: "Sheep",
        23: "Sofa",
        24: "Train",
        25: "Tvmonitor",
        26: "Pedestrian",
        27: "Tricycle",
        28: "Awning-tricycle"
    }
    origin_classes = [value for k,value in origin_dict.items()]
    Extra_classes = down_load_class
    
    dataset = foz.load_zoo_dataset(
        "open-images-v6",
        split = "validation",
        max_samples=max_samples,
        seed=seed,
        shuffle=shuffle,
        label_types=["detections"],
        classes = Extra_classes,
        #only_matching=True,
    )
    new_dict = origin_dict
    for c in Extra_classes:
        if c not in origin_classes:
            new_dict[len(new_dict)] = c
    new_classes = [value for k,value in new_dict.items()]
    print(len(new_classes))

    os.makedirs('../openimages/train', exist_ok=True)
    os.makedirs('../openimages/val', exist_ok=True)
    os.makedirs('../openimages/test', exist_ok=True)

    train_set = dataset.take(int(0.7 * len(dataset)), seed=42)
    val_set = dataset.exclude(train_set).take(int(0.2 * len(dataset)), seed=42)
    test_set = dataset.exclude(train_set).exclude(val_set)

    
    train_set.export(
        export_dir="../openimages/train",
        dataset_type=fo.types.YOLOv5Dataset,
        label_field="ground_truth",  # 标签字段名
        classes=new_classes,
    )
    val_set.export(
        export_dir="../openimages/val",
        dataset_type=fo.types.YOLOv5Dataset,
        label_field="ground_truth",  # 标签字段名
        classes=new_classes,
    )
    
    test_set.export(
        export_dir="../openimages/test",
        dataset_type=fo.types.YOLOv5Dataset,
        label_field="ground_truth",  # 标签字段名
        classes=new_classes,
    )



def visdrone2yolo(dir, convert_dict):
  from PIL import Image
  from tqdm import tqdm

  def convert_box(size, box):
      # Convert VisDrone box to YOLO xywh box
      dw = 1. / size[0]
      dh = 1. / size[1]
      return (box[0] + box[2] / 2) * dw, (box[1] + box[3] / 2) * dh, box[2] * dw, box[3] * dh

  (dir / 'labels').mkdir(parents=True, exist_ok=True)  # make labels directory
  pbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}')
  for f in pbar:
      img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).size
      lines = []
      with open(f, 'r') as file:  # read annotation.txt
          for row in [x.split(',') for x in file.read().strip().splitlines()]:
              if row[4] == '0':  # VisDrone 'ignored regions' class 0
                  continue
              cls = convert_dict[int(row[5]) - 1]
              box = convert_box(img_size, tuple(map(int, row[:4])))
              lines.append(f"{cls} {' '.join(f'{x:.6f}' for x in box)}\n")
              with open(str(f).replace(os.sep + 'annotations' + os.sep, os.sep + 'labels' + os.sep), 'w') as fl:
                  fl.writelines(lines)  # write label.txt
                  
def download_VisDrone(VisDrone_path = '../VisDrone'):
    convert_dict = {
        # 这个是Vis的类别转到整个的类别的字典
        0: 26,#pedestrian
        1: 4,#people
        2: 9,#bicycle
        3: 0,#car
        4: 1,#van
        5: 2,#truck
        6: 27,#tricycle --三轮车
        7: 28,#awning-tricycle -- 带遮阳棚的三轮车
        8: 13,#bus
        9: 20,#motor
    }
    dir = Path(VisDrone_path)  # dataset root dir
    urls = ['https://github.com/ultralytics/assets/releases/download/v0.0.0/VisDrone2019-DET-train.zip',
          'https://github.com/ultralytics/assets/releases/download/v0.0.0/VisDrone2019-DET-val.zip',
          'https://github.com/ultralytics/assets/releases/download/v0.0.0/VisDrone2019-DET-test-dev.zip',
          'https://github.com/ultralytics/assets/releases/download/v0.0.0/VisDrone2019-DET-test-challenge.zip']
    download(urls, dir=dir, curl=True, threads=4)

    # Convert
    for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':
      visdrone2yolo(dir / d, convert_dict)  # convert VisDrone annotations to YOLO labels

    



